
import os
import sys
import json
import time
import shutil
import requests
import numpy as np
import soundfile as sf
from tqdm import tqdm

import utils

sys.path.append('./src/')
sys.path.append('./src/audiohandler/speechbrainlib/')
sys.path.append('./src/separationhandler/asteroid/')

from tools.courseinfo import course_info
from tools.getscore import get_final_score
from facehandler.task1handler import Task1Handler
from audiohandler.task2handler import Task2Handler
from separationhandler.task3handler import Task3Handler


def test_task1(video_path):

    print('*'*40)
    print('Start processing Task1 @ ' + time.strftime('%Y-%m-%d %H:%M:%S', time.localtime()))
    print('-'*5)
    start_time = time.time()

    print('[Stage 1/4] Init Dataset & Model')

    if not os.path.exists('./dataset/task1/model/'):
        os.makedirs('./dataset/task1/model/')

    if os.path.exists("./dataset/task1/model/shape_predictor_68_face_landmarks.dat"):
        print('  Model 1: (shape_predictor_68_face_landmarks.dat)')
        print('  Model 2: (dlib_face_recognition_resnet_model_v1.dat)')
    else:
        # 训练好的模型下载
        url1 = "https://cloud.tsinghua.edu.cn/f/245647e72734495d8f63/?dl=1" # 关键点定位模型
        url2 = "https://cloud.tsinghua.edu.cn/f/cf85d429bdf14d78a7ea/?dl=1" # 面部识别模型

        r1 = requests.get(url1)
        r2 = requests.get(url2)

        with open ("./dataset/task1/model/shape_predictor_68_face_landmarks.dat","wb") as f1:
            f1.write(r1.content)
            f1.close()
        with open ("./dataset/task1/model/dlib_face_recognition_resnet_model_v1.dat","wb") as f2:
            f2.write(r2.content)
            f2.close()
        
        print('  Download Model 1: (shape_predictor_68_face_landmarks.dat)')
        print('  Download Model 2: (dlib_face_recognition_resnet_model_v1.dat)')

    if not os.path.exists('./dataset/task1/featureDB'):
        url = 'https://cloud.tsinghua.edu.cn/f/b0151fa06b6a4b35b325/?dl=1'
        r = requests.get(url)
        with open ('feature.zip',"wb") as f:
            f.write(r.content)
            f.close()
        print('  Download Pretrained Models: (feature.zip)')
        os.system('unzip -d ./dataset/task1 feature.zip > /dev/null 2>&1 && rm -rf feature.zip')
    
    task1Handler = Task1Handler()
    result_dict = task1Handler.returndic(video_path)
    
    end_time = time.time()
    print("Processing Task 1 Done, Using Time %.2f s" % (end_time - start_time))
    print('[Stage 4/4] Evaluating')

    return result_dict

def test_task2(wav_path):

    print('\n'+'*'*40)
    print('Start processing Task2 @ ' + time.strftime('%Y-%m-%d %H:%M:%S', time.localtime()))
    print('-'*5)
    start_time = time.time()

    result_dict = {}
    if not os.path.exists('./pretrained_models'):
        url = 'https://cloud.tsinghua.edu.cn/f/256185d279b74871a2c3/?dl=1'
        r = requests.get(url)
        with open ('pretrained_models.zip',"wb") as f:
            f.write(r.content)
            f.close()
        print('  Download Pretrained Models: (pretrained_models.zip)')
        os.system('unzip pretrained_models.zip > /dev/null 2>&1 && rm -rf pretrained_models.zip')
    audio_handler = Task2Handler(2)
    audio_handler.test_init()

    print('[Stage 5/6] Verify Audios')
    for file_idx in tqdm(range(len(os.listdir(wav_path))), ncols=70):
        file_name = os.listdir(wav_path)[file_idx]
        person_id, score_list = audio_handler.test_single_audio(os.path.join(wav_path, file_name))
        result_dict[file_name]=utils.ID_dict[person_id]
    
    end_time = time.time()
    print("Processing Task 2 Done, Using Time %.2f s" % (end_time - start_time))
    print('[Stage 6/6] Evaluating')

    return result_dict

def test_task3(video_path,result_path):
    
    print('\n'+'*'*40)
    print('Start processing Task3 @ ' + time.strftime('%Y-%m-%d %H:%M:%S', time.localtime()))
    print('-'*5)
    if os.path.isdir(result_path):
        # print('warning: using existed path as result_path')
        pass
    else:
        os.mkdir(result_path)
    start_time = time.time()

    target_path='./dataset/test_offline/task3_result'
    if os.path.exists(target_path):
        os.system('rm -rf '+target_path)
    new_temp_audio_path='./src/separationhandler/data/result_new'

    separation_handler = Task3Handler(
        model_name='JorisCos/ConvTasNet_Libri3Mix_sepclean_8k',
        video_path=video_path,
        audio_path='./dataset/task3_audio',
        face_label_path='./src/separationhandler/data/facedata.json',
        match_table_path='./src/separationhandler/tools/matchtable.csv',
        sep_list_path='./src/separationhandler/data/test_demo.csv',
        temp_audio_path='./src/separationhandler/data/result',
        new_temp_audio_path=new_temp_audio_path
    )

    print('[Stage 1/8] Check Path & Model')
    separation_handler.check_path()
    if not os.path.exists('./src/separationhandler/data/facedata.json'):
        print('  Make Face Label: (' + './src/separationhandler/data/facedata.json' + ')')
        os.system('python ./src/multifacehandler/task3_recognition.py')
    else:
        print('  Face Label (' + './src/separationhandler/data/facedata.json' + ')')
    separation_handler.process_video()

    print('[Stage 2/8] Load Model')
    separation_handler.load_model()

    print('[Stage 3/8] Separate Audios')
    separation_handler.separate_3_audio()

    print('[Stage 4/8] Reprocessing - Rename Audios')
    separation_handler.rename_audio()

    print('[Stage 5/8] Reprocessing - Get Match Table')
    separation_handler.get_match_table()

    print('[Stage 6/8] Reprocessing - Add Position Info')
    separation_handler.add_audio_position()

    shutil.copytree(new_temp_audio_path, target_path)

    print('[Stage 7/8] Reprocessing - Resample Audios')
    os.system('cd tools && python resample.py')
    
    end_time = time.time()
    print("Processing Task 3 Done, Using Time %.2f s" % (end_time - start_time))
    print('[Stage 8/8] Evaluating')

if __name__=='__main__':

    course_info()
    
    ## testing task1
    with open('./dataset/test_offline/task1_gt.json','r') as f:
        task1_gt = json.load(f)
    task1_pred = test_task1('./dataset/test_offline/task1')
    task1_acc = utils.calc_accuracy(task1_gt,task1_pred)
    print('---------------------------------------------------')
    print('|                Task1 Eval Result                |')
    print('---------------------------------------------------')
    print('*[indicator] accuracy for task1 is:',task1_acc)
    print('---------------------------------------------------')

    ## testing task2
    with open('./dataset/test_offline/task2_gt.json','r') as f:
        task2_gt = json.load(f)
    task2_pred = test_task2('./dataset/test_offline/task2')
    task2_acc = utils.calc_accuracy(task2_gt, task2_pred)
    print('---------------------------------------------------')
    print('|                Task2 Eval Result                |')
    print('---------------------------------------------------')
    print('*[indicator] accuracy for task2 is:', task2_acc)
    print('---------------------------------------------------')

    ## testing task3
    test_task3('./dataset/test_offline/task3','./dataset/test_offline/task3_estimate')
    print('---------------------------------------------------')
    print('|                Task3 Eval Result                |')
    print('---------------------------------------------------')
    task3_SISDR_blind = utils.calc_SISDR('./dataset/test_offline/task3_gt', './dataset/test_offline/task3_estimate', permutaion=True)  # 盲分离
    print('*[indicator] strength-averaged SISDR_blind for task3 is: %.3f' % task3_SISDR_blind)
    task3_SISDR_match = utils.calc_SISDR('./dataset/test_offline/task3_gt', './dataset/test_offline/task3_estimate', permutaion=False) # 定位分离
    print('*[indicator] strength-averaged SISDR_match for task3 is: %.3f' % task3_SISDR_match)
    print('---------------------------------------------------')
    print()
    print('---------------------------------------------------')
    print('|                   Final Score                   |')
    print('---------------------------------------------------')
    print('*[final score] ', end='')
    print('%.3f / 50' % get_final_score(task1_acc, task2_acc, task3_SISDR_blind, task3_SISDR_match))
